import os
import shutil  # 用于文件/目录的复制、删除等操作


class Recorder(object):
    """
    实验记录器类：用于管理实验过程中的文件存储、日志记录、代码备份等功能
    核心作用是规范化实验结果的保存，便于后续复现和追溯
    """

    def __init__(self, snapshot_pref, ignore_folder):
        """
        初始化记录器
        Args:
            snapshot_pref: 实验结果保存的根目录路径（字符串）
            ignore_folder: 代码备份时需要忽略的目录名称（字符串）
        """
        # 若保存根目录不存在，则创建该目录
        if not os.path.isdir(snapshot_pref):
            os.mkdir(snapshot_pref)

        # 定义各类文件/目录的保存路径
        self.save_path = snapshot_pref  # 实验结果根目录
        self.log_file = self.save_path + "log.txt"  # 实验日志文件路径
        self.readme = self.save_path + "README.md"  # 实验说明文档路径
        self.opt_file = self.save_path + "opt.log"  # 实验参数配置文件路径
        self.code_path = os.path.join(self.save_path, "code/")  # 代码备份目录路径

        # 若README文件已存在，则删除（避免历史内容干扰）
        if os.path.isfile(self.readme):
            os.remove(self.readme)

        # 若代码备份目录不存在，则创建
        if not os.path.isdir(self.code_path):
            os.mkdir(self.code_path)

        # 备份当前项目代码到指定目录（忽略指定文件夹）
        self.copy_code(dst=self.code_path, ignore_folder=ignore_folder)

        # 打印实验结果保存路径，提示用户
        print("\n======> Result will be saved at: ", self.save_path)

    def copy_code(self, src="./", dst="./code/", ignore_folder='Exps'):
        """
        代码备份方法：将源目录下的代码文件复制到目标目录，支持忽略指定文件夹
        Args:
            src: 源代码目录路径（默认当前目录）
            dst: 代码备份目标目录路径
            ignore_folder: 需要忽略的目录名称（默认'Exps'，通常为实验结果目录）
        """
        import uuid  # 用于生成唯一标识符，避免目录重名冲突

        # 若目标目录已存在，在目录名后添加唯一标识符（避免覆盖已有备份）
        if os.path.isdir(dst):
            dst = "/".join(dst.split("/")) + "code_" + str(uuid.uuid4()) + "/"

        # 存储所有需要备份的文件绝对路径
        file_abs_list = []
        src_abs = os.path.abspath(src)  # 获取源目录的绝对路径

        # 遍历源目录下的所有文件和子目录
        for root, dirs, files in os.walk(src_abs):
            # 忽略包含指定文件夹的路径（避免备份实验结果等无关目录）
            if ignore_folder not in root:
                for name in files:
                    # 记录符合条件的文件绝对路径
                    file_abs_list.append(root + "/" + name)

        # 遍历需要备份的文件，执行复制操作
        for file_abs in file_abs_list:
            file_split = file_abs.split("/")[-1].split('.')  # 拆分文件名和后缀
            # 筛选条件：文件大小小于10MB，且不是.pyc编译文件（避免备份大文件和无用编译产物）
            if os.path.getsize(file_abs) / 1024 / 1024 < 10 and not file_split[-1] == "pyc":
                src_file = file_abs  # 源文件路径
                # 构造目标文件路径（保持与源目录一致的目录结构）
                dst_file = dst + file_abs.replace(src_abs, "")
                # 若目标文件所在目录不存在，则创建
                if not os.path.exists(os.path.dirname(dst_file)):
                    os.makedirs(os.path.dirname(dst_file))
                # 复制文件到目标路径（尝试两次，避免偶发错误）
                try:
                    shutil.copyfile(src=src_file, dst=dst_file)
                except:
                    print("copy file error")  # 复制失败时打印错误提示

    def writeopt(self, opt):
        """
        写入实验参数配置：将实验的参数对象保存到opt.log文件
        Args:
            opt: 参数对象（通常为argparse解析后的对象，包含各类实验配置）
        """
        with open(self.opt_file, "w") as f:
            # 遍历参数对象的属性，按"键: 值"格式写入文件
            for k, v in opt.__dict__.items():
                f.write(str(k) + ": " + str(v) + "\n")

    def writelog(self, input_data):
        """
        写入实验日志：将实验过程中的日志信息追加到log.txt文件
        Args:
            input_data: 待写入的日志内容（字符串或可转换为字符串的对象）
        """
        txt_file = open(self.log_file, 'a+')  # 以追加模式打开日志文件
        txt_file.write(str(input_data) + "\n")  # 写入内容并换行
        txt_file.close()  # 关闭文件

    def writereadme(self, input_data):
        """
        写入实验说明：将实验的关键说明信息追加到README.md文件
        Args:
            input_data: 待写入的说明内容（支持Markdown格式）
        """
        txt_file = open(self.readme, 'a+')  # 以追加模式打开README文件
        txt_file.write(str(input_data) + "\n")  # 写入内容并换行
        txt_file.close()  # 关闭文件

    def gennetwork(self, var):
        """
        生成网络结构图表（当前方法未实现完整逻辑，仅定义接口）
        Args:
            var: 网络相关变量（推测为模型对象或计算图变量）
        """
        self.graph.draw(var=var)  # 依赖未定义的graph属性，需结合其他代码补充

    def savenetwork(self):
        """
        保存网络结构图表（当前方法未实现完整逻辑，仅定义接口）
        将网络结构保存为SVG格式文件到实验结果目录
        """
        self.graph.save(file_name=self.save_path + "network.svg")  # 依赖未定义的graph属性

    """
    以下方法为注释掉的备用功能（权重记录与可视化），当前未启用：
    - writeweights: 记录网络层权重数据到日志文件
    - drawhist: 基于权重日志绘制权重分布直方图并保存
    """
    """def writeweights(self, input_data, block_id, layer_id, epoch_id):
        txt_path = self.weight_folder + "conv_weight_" + str(epoch_id) + ".log"
        txt_file = open(txt_path, 'a+')
        write_str = "%d\t%d\t%d\t" % (epoch_id, block_id, layer_id)
        for x in input_data:
            write_str += str(x) + "\t"
        txt_file.write(write_str+"\n")

    def drawhist(self):
        drawer = DrawHistogram(txt_folder=self.weight_folder, fig_folder=self.weight_fig_folder)
        drawer.draw()"""